﻿using System;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Assertions;
using UnityEngine.Events;
using UnityEngine.Rendering;



namespace XYRendering.Vegetation
{

  public enum VGBatchingOption
  {
    Batching_511,
    Batching_1023,
  }


  #region Batches

  public class VGInstancingBatch
  {
    public Mesh mesh;
    public Material material;
    public int numInstances;
    public readonly int batchSize;
    public readonly Matrix4x4[] instancesMatrixs;
    public readonly MaterialPropertyBlock mpb;
    
    private readonly Vector3[] _instancesWorldPositions;
    
    private static Vector4[] s_OcclusionDatas = null;
    private static SphericalHarmonicsL2[] s_LightProbesDatas = null;

        
    public bool isFull => numInstances >= batchSize;
    
    public bool canExecute => mesh != null && material != null && numInstances > 0;
    
    public VGInstancingBatch(int batchSz)
    {
      mesh = null;
      material = null;
      numInstances = 0;
      batchSize = batchSz;
      instancesMatrixs = new Matrix4x4[batchSz];
      mpb = new MaterialPropertyBlock();
     
      _instancesWorldPositions = new Vector3[batchSz]; 
     
      if (s_OcclusionDatas == null)
        s_OcclusionDatas = new Vector4[batchSz];
      
      if (s_LightProbesDatas == null)
        s_LightProbesDatas = new SphericalHarmonicsL2[batchSz];
      
      // call here for memory allocation mpb memory
      mpb.CopySHCoefficientArraysFrom(s_LightProbesDatas);
      mpb.CopyProbeOcclusionArrayFrom(s_OcclusionDatas);
    }
    
    public virtual bool AddInstance(Matrix4x4 transform)
    {
      if (isFull)
        return false;
        
      instancesMatrixs[numInstances] = transform;
      _instancesWorldPositions[numInstances] = transform.GetColumn(3);
      
      numInstances++;
      
      return true;
    }

    public virtual void UpdateOcclusionAndLightProbesData()
    {
      LightProbes.CalculateInterpolatedLightAndOcclusionProbes(_instancesWorldPositions, s_LightProbesDatas, s_OcclusionDatas);
      mpb.CopyProbeOcclusionArrayFrom(s_OcclusionDatas);
      mpb.CopySHCoefficientArraysFrom(s_LightProbesDatas);
    }
    
    public virtual void Release()
    {
      mesh = null;
      material = null;
      numInstances = 0;
    }
    
  }

  
  public class VGInstancingBatch511 : VGInstancingBatch
  {
    public VGInstancingBatch511() : base(511) { }
  }
  
  
  public class VGInstancingBatch1023 : VGInstancingBatch
  {
    public VGInstancingBatch1023() : base(1023) {}
  }

  #endregion



  #region Batch Utils

  public class VGBatchUtils
  {
  
    // ----------------------------------------------------- Batch Pool -----------------------------------------------
    private interface VGBatchPool
    {
      public VGInstancingBatch Get();
      public void Release(VGInstancingBatch batch);
    }
    

    private class VGBatchPool511 : VGBatchPool 
    {
    
      ObjectPool<VGInstancingBatch511> _pool;

      public VGBatchPool511(UnityAction<VGInstancingBatch511> getOp, UnityAction<VGInstancingBatch511> releaseOp)
      {
        _pool = new ObjectPool<VGInstancingBatch511>(getOp, releaseOp);
      }
    
      public VGInstancingBatch Get()
      {
        return _pool.Get();
      }

      public void Release(VGInstancingBatch obj)
      {
        _pool.Release((VGInstancingBatch511)obj);
      }
    }


    private class VGBatchPool1023 : VGBatchPool
    {
      ObjectPool<VGInstancingBatch1023> _pool;

      public VGBatchPool1023(UnityAction<VGInstancingBatch1023> getOp, UnityAction<VGInstancingBatch1023> releaseOp)
      {
        _pool = new ObjectPool<VGInstancingBatch1023>(getOp, releaseOp);
      }
    
      public VGInstancingBatch Get()
      {
        return _pool.Get();
      }

      public void Release(VGInstancingBatch obj)
      {
        _pool.Release((VGInstancingBatch1023)obj);
      }
    }
    
    
    // ----------------------------------------------------- Batch Pool -----------------------------------------------

    public readonly  VGBatchingOption batchingOption;
    public readonly  int maxNumBatches;
    
    public int numInstances { get; private set; }
    public int numBatches { get; private set; }
    
      
    VGInstancingBatch[] _allBatches;
    
    Dictionary<UInt64, int> _avalibleBatches;


    VGBatchPool _batchPool;
    
    public VGInstancingBatch this[int idx] => _allBatches[idx];

    public VGBatchUtils(VGBatchingOption batchingOption, int maxNumBatches = 2048)
    {
      this.batchingOption = batchingOption;
      this.maxNumBatches = maxNumBatches;
      
      numInstances = 0;
      numBatches = 0;
      
      _avalibleBatches = new Dictionary<ulong, int>();
      
      _allBatches = new VGInstancingBatch[maxNumBatches];
      
      _batchPool = batchingOption == VGBatchingOption.Batching_511 ? (VGBatchPool)new VGBatchPool511(null, (batch) => batch.Release())
                                   : (VGBatchPool)new VGBatchPool1023(null, (batch) => batch.Release());
      
    }
    
    
    public static UInt64 ComputeBatchID(Mesh mesh, Material material)
    {
        return (UInt64)(uint)mesh.GetInstanceID() << 32 | (uint)material.GetInstanceID();
    }

    public void Clear()
    {
      for (int i = 0; i < numBatches; i++)
      {
        _batchPool.Release(_allBatches[i]);
        _allBatches[i] = null;
      }
      numBatches = 0;
      numInstances = 0;
      _avalibleBatches.Clear();
    }

    public bool CommitInstance(Mesh mesh, Material material, Matrix4x4 transform)
    {
        UInt64 batchId = ComputeBatchID(mesh, material);
        VGInstancingBatch batch = GetAvalibleBatch(batchId, mesh, material);
        bool success=  batch.AddInstance(transform);
        
        if (success)
          numInstances++;
          
        return success;
    }

    VGInstancingBatch GetAvalibleBatch(UInt64 batchId, Mesh mesh, Material material)
    {
      VGInstancingBatch freeBatch = null;
      if (_avalibleBatches.ContainsKey(batchId))
        freeBatch = _allBatches[_avalibleBatches[batchId]];
      
      // double check, if no exist avalible batch or  current batch is full, rank a new one from pool
      if (freeBatch == null || freeBatch.isFull)
      {
        if (numBatches >= maxNumBatches) // we reach maximum batch limit, can't get any more
          throw new OutOfMemoryException();
          
        freeBatch = _batchPool.Get();
        freeBatch.mesh = mesh;
        freeBatch.material = material;
        
        _allBatches[numBatches] = freeBatch;
        _avalibleBatches[batchId] = numBatches;
        
        numBatches++;
      }

#if UNITY_EDITOR
      Assert.IsNotNull(freeBatch);
      Assert.IsTrue(freeBatch.numInstances < freeBatch.batchSize);
#endif
      
      return freeBatch;
    }
    
  }

  #endregion
  
  
    
}